import sys
import os
import codecs
import re
import json
import argparse
import subprocess
import shutil
import tempfile
import time
import datetime
import glob
import urllib.request
import urllib.parse
import urllib.error
import urllib.request
import urllib.error
import urllib.parse
import urllib.parse
import urllib.error
import urllib.request
import urllib
import multiprocessing
import numpy as np
import psutil
from typing import List, Dict, Tuple, Union, Optional, Any
from tqdm import tqdm

# cat /home/gaoxinglong/env/data/south3k/gudo/0000021/00000210015.wav | wav-reverberate --shift-output=true --impulse-response="sox RIRS_NOISES/real_rirs_isotropic_noises/RWCP_type1_rir_circle_e1c_imp160.wav -r 16000 -t wav - |"  - - |


def get_real_filename(cmdline):
    if len(cmdline) == 0:
        return None
    line_tokens = re.split(r'\s+', cmdline)
    if len(line_tokens) > 2:
        # return full path, realpath, and extended filename
        return os.path.basename(os.path.realpath(line_tokens[1])), os.path.realpath(line_tokens[1]), os.path.splitext(os.path.realpath(line_tokens[1]))[1]


def conv_cmd_to_simple(cmdline):
    if len(cmdline) == 0:
        return None
    line_tokens = re.split(r'\s+', cmdline)
    if len(line_tokens) > 2:
        return " ".join(line_tokens[3:13])


def exec_shell_cmd_to_create_real_wav_file(cmd_list: List[str], out_wav_list: multiprocessing.Queue):
    """
    Execute shell command to create a real wav file
    :param cmd_list:
    :param out_wav_file:
    :return:
    """
    local_scps = []
    for cmd in tqdm(cmd_list):
        line_tokens = re.split(r'\s+', cmd, maxsplit=1)
        if len(line_tokens) == 2:
            one_cmd = line_tokens[1]
            filename, fullfilename, ext_name = get_real_filename(one_cmd)
            filename = re.sub(ext_name, '', filename)
            fullfilename = re.sub(ext_name, '', fullfilename)
            outfullfilename = '{}{}'.format(fullfilename, '_rir.wav')
            simple_cmd = conv_cmd_to_simple(one_cmd)
            full_cmd = '{} {} {}'.format(
                simple_cmd, fullfilename + '.wav', outfullfilename)

            local_scps.append('{}\t{}'.format(line_tokens[0], outfullfilename))
            os.system(command=full_cmd + "> /dev/null 2>&1")
    out_wav_list.put(local_scps)


def parallel_process(input_cmd_list: str, raw_rir_root: str, num_jobs: int):
    """
    Parallel process
    :param input_cmd_list:
    :param out_wav_file_list:
    :param num_jobs:
    :return:
    """

    with codecs.open(input_cmd_list, 'r', encoding='utf-8') as f:
        cmd_list = f.readlines()
        cmd_list = [re.sub(r'RIRS_NOISES', raw_rir_root, cmd)
                    for cmd in cmd_list]
        num_jobs_cpu = len(cmd_list) // num_jobs + 1
        with multiprocessing.Manager() as manager:
            pool = multiprocessing.Pool(processes=num_jobs)
            q = manager.Queue()
            for i in range(num_jobs):
                tasks_every_cpu = cmd_list[i *
                                           num_jobs_cpu:min(len(cmd_list), (i + 1) * num_jobs_cpu)]
                pool.apply_async(
                    exec_shell_cmd_to_create_real_wav_file, args=(tasks_every_cpu, q))
            pool.close()
            pool.join()
            with codecs.open(input_cmd_list + '.rir.scp', 'w', encoding='utf-8') as f:
                while not q.empty():
                    for line in q.get():
                        f.write(line + '\n')


if __name__ == "__main__":
    raw_rir_root = '/home/share_ssd_data/env/weimeng/datasets/SLR28-Room-Impulse-Response-and-Noise-Database'
    cmd_wav_scp = sys.argv[1]
    num_cpus = psutil.cpu_count(logical=True)
    parallel_process(cmd_wav_scp, raw_rir_root, num_jobs=num_cpus)
    print('Done!')
